package com.lty.redis;

import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶算法限流
 * @author lty
 */
@Slf4j
@Component
public class RedisRaterLimiter {

    @Resource
    private RedissonClient redisson;

    /**
     * 限流前缀
     */
    private String LIMIT_PRE="LIMIT:";

    /**
     * 基于Redis令牌桶算法
     * @param name         限流标识(基本为方法名+ip作为唯一标识)
     * @param rate         限制的数量 速率
     * @param rateInterval 单位时间内（毫秒）
     * @return
     */
    public Boolean acquireByRedis(String name, Long rate, Long rateInterval) {
        boolean getToken;
        try {
            RRateLimiter rateLimiter = redisson.getRateLimiter(this.LIMIT_PRE + name);
            rateLimiter.trySetRate(RateType.OVERALL, rate, rateInterval, RateIntervalUnit.MILLISECONDS);
            getToken = rateLimiter.tryAcquire();
            rateLimiter.expireAsync(rateInterval * 2, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            getToken = true;
        }
        return getToken;
    }
}
